import os
import PIL.Image as Image
import numpy as np
import torch
import torch.utils.data as data
import torchvision.transforms as transforms
from utils.config import *
from utils.io_helper import read

# root = 'G:/USVInland/Stereo Matching/Low_Res_640_320' # Windows
# root = '/home/ubuntu/zhouyiqing/USVInland/Stereo Matching/Low_Res_640_320' # Ubuntu

# USVInland数据集
class USVInland(data.Dataset):
    def __init__(self, root, transform, seg):
        self.root = root
        self.transform = transform

        # 拼接左右图像文件夹路径
        if seg: # 分割水体
            left_path = os.path.join(root, 'Left_Img_Rectified_Seg_rm')
            right_path = os.path.join(root, 'Right_Img_Rectified_Seg_rm')
            disp_path = os.path.join(root, 'Disp_Map_Seg_rm')
        else: # 未分割
            left_path = os.path.join(root, 'Left_Img_Rectified')
            right_path = os.path.join(root, 'Right_Img_Rectified')
            disp_path = os.path.join(root, 'Disp_Map')

        left_imgs = list()
        right_imgs = list()
        disp_imgs = list()

        # 拼接图片文件名，并加入文件名list
        left_imgs = list(os.path.join(left_path, img) for img in os.listdir(left_path))
        right_imgs = list(os.path.join(right_path, img) for img in os.listdir(right_path))
        disp_imgs = list(os.path.join(disp_path, img) for img in os.listdir(disp_path))

        self.left_imgs = left_imgs
        self.right_imgs = right_imgs
        self.disp_imgs = disp_imgs

    def __getitem__(self, index):
        # 读取RGB图像
        left_img = Image.open(self.left_imgs[index]).convert('RGB')
        right_img = Image.open(self.right_imgs[index]).convert('RGB')
        if system == 'Windows': img_name = self.left_imgs[index].split('\\')[-1]
        elif system == 'Ubuntu': img_name = self.left_imgs[index].split('/')[-1]

        # 读取视差灰度图像
        disp_img = read(self.disp_imgs[index])
        disp_img = Image.fromarray(disp_img.astype(np.uint8), mode='L')

        # 图像预处理
        if self.transform:
            seed = torch.random.seed()
            torch.random.manual_seed(seed)
            left_img = self.transform(left_img)
            torch.random.manual_seed(seed)
            right_img = self.transform(right_img)

            disp_transform = transforms.Compose([transforms.ToTensor()])
            torch.random.manual_seed(seed)
            disp_img = disp_transform(disp_img)[0, :] * 255.0 # 灰度图ToTensor自动归一化到[0, 1]，还原到[0, 255]

            # 张量转图片显示
            # transforms.ToPILImage()(left_img).convert('RGB').show()
            # transforms.ToPILImage()(right_img).convert('RGB').show()
            # transforms.ToPILImage()(disp_img).convert('L').show()

            # 数据保存到字典
            data = {'left': left_img, 'right': right_img, 'gt': disp_img, 'name': img_name}

        return data

    def __len__(self):
        return len(self.left_imgs)
